class SimulationParameters:
	
	def __init__(self, iPopulation, iDemography):
		self.myPopulation = iPopulation
		self.myDemography = iDemography
		
	def getMSString_Null(self, numSims=None):
		s = ''
		s += self.myPopulation.getMSString_Null(numSims)
		s += self.myDemography.getMSString()
		return s
	
	def getMSString_Alternative(self, altDemography, numAltSims=None):
		s = ''
		s += self.myPopulation.getMSString_Alternative(numAltSims)
		s += altDemography.getMSString()
		return s
		
	def getNumChromosomes(self):
		return self.myPopulation.numChromosomes
	
	def getNumSimulations(self):
		return self.myPopulation.numSimulations

	def getSegSites(self):
		return self.myPopulation.numSegregatingSites

	def isValidOperator(self, iLinOp):
		n = iLinOp.size()
		return (n == self.myPopulation.numOfChromosomes)
	
	def __str__(self):
		s = str(self.myPopulation) + '\n'
#		s += str(self.myDemography)
		return s



'''
These are parameters that shouldn't change between different demographies
Updated Jan 12 2010
'''
class Population:

	numChromosomes = 0
	numSegregatingSites = 0

	rateMutation = 0.0
	rateRecombination = 0.0
	sequenceLength = 0

	N = 0
	numSimulations = 0
	numPowerSimulations = 0
	
	sequenceLength = 0
		

	def getTheta(self):
		theta = 4.0 * self.N * self.rateMutation * self.sequenceLength
		return theta

	def getRho(self):
		
		# float equality for zero
		if self.rateMutation < 1e-12:
			return 0.0
		
		rho = 4.0 * self.N * self.rateRecombination * self.sequenceLength
		return rho

	def getMSString_Null(self, numSims=None):
		
		nSims = self.numSimulations
		
		if numSims is not None:
			nSims = numSims
			
		if self.numSegregatingSites < 1:
			s = ' %d %d -t %f' % (self.numChromosomes, nSims, self.getTheta() )
		else:
			s = ' %d %d -s %d' % (self.numChromosomes, nSims, self.numSegregatingSites )
			
		r = self.getRho()
		
		# float inequality for zero
		if r > 1e-12:
			s += ' -r %f %d' % ( r, self.sequenceLength )
		
		return s

	def getMSString_Alternative(self, numAltSims=None):
		
		nSims = self.numPowerSimulations
		
		if numAltSims is not None:
			nSims = numAltSims
			
		if self.numSegregatingSites < 1:
			s = ' %d %d -t %f' % (self.numChromosomes, nSims, self.getTheta() )
		else:
			s = ' %d %d -s %d' % (self.numChromosomes, nSims, self.numSegregatingSites )
			
		r = self.getRho()
		
		# float inequality for zero
		if r > 1e-12:
			s += ' -r %f %d' % ( r, self.sequenceLength )
		
		return s

	def __str__(self):
		value  = 'nChr:\t ' + str(self.numChromosomes)
		value += '\t Ne:\t ' + str(self.N)
		value += '\t L:\t ' + str(self.sequenceLength)
		value += '\t mu:\t ' + str(self.rateMutation)
		value += '\t r:\t ' + str(self.rateRecombination) 
		
		return value

'''
A Demography is made up of epochs
Updated Jan 12 2010
'''
class Demography:

	def __init__(self, N):
		self.N = N
		self.epochs = []

	def addEpoch(self, alpha, tau, isConstant):
		ep = None
		
		if isConstant:
			ep = C_Epoch( alpha, tau, self.N )
		else:
			ep = G_Epoch( alpha, tau, self.N )

		self.epochs.append( ep )
		return ep

	def getMSString(self):
		s = ''
		for ep in self.epochs:
			s += ep.getMSString()
		return s

	def __str__(self):
		s = "N;%d;" % (self.N)
		
		for ep in self.epochs:
			s += str(ep)
		
		return s

'''
An Epoch has a tau and alpha parameter and type (C/G)
The N is the scaling factor
The two types of Epochs are Constant and Growth
Updated Jan 12, 2010
'''
class Epoch:
	
	def __init__( self, alpha, tau, N, isConstant ):
		''' Tau is in generation time, alpha is what it is '''
		
		self.N = N
		
		self.tau = 1.0 * tau / (4.0 * N)
		self.alpha = alpha
		
		self.isConstant = isConstant
		
		
		self.itrParam = None
		self.itrIncrement = None
		self.itrStop= None

	def setIterationParameters(self, iParameter, iStart, iIncrement, iStop, iLog = False):
		self.itrParam = iParameter
		
		if iParameter == 'alpha':
			self.alpha = iStart
		else: # 'tau'
			self.tau = iStart
			
		self.itrIncrement = iIncrement
		self.itrStop = iStop
		self.itrLog = iLog

	def next(self):
		if self.itrParam == 'alpha':
			if self.itrLog:
				if (self.alpha * self.itrIncrement > self.itrStop):
					raise StopIteration
				self.alpha *= self.itrIncrement
			else:
				if (self.alpha + self.itrIncrement > self.itrStop):
					raise StopIteration
				self.alpha += self.itrIncrement
				
			return self

		elif self.itrParam == 'tau':
			if self.itrLog:
				if self.tau * self.itrIncrement > self.itrStop:
					raise StopIteration
				self.tau *= self.itrIncrement
			else:
				if self.tau + self.itrIncrement > self.itrStop:
					raise StopIteration
				self.tau += self.itrIncrement
			return self

		else:
			raise StopIteration

	def __iter__(self):
		return self
	
	def __str__(self):
		value = ""
		
		if self.isConstant:
			value += "C;"
		else:
			value += "G;"
			
		value  += "%f;%f" % (self.alpha, self.tau)
		
		return value


'''
Basically the Wright-Fisher model
Updated Jan 12, 2010
'''
class C_Epoch( Epoch ):
	
	def __init__(self, alpha, tau, N ):
		Epoch.__init__(self, alpha, tau, N, True)

	def getMSString( self ):

		rValue = ' '
		
		# float equality for zero
		if self.tau < 1e-12:
			return rValue
		
		self.popFraction = 1.0 * self.alpha / self.N
		rValue = ' -eN %f %f' % (self.tau, self.popFraction)
		
		return rValue

'''
Model for a population that's continuously expanding
Updated Jan 12, 2010
'''
class G_Epoch( Epoch ):
	
	def __init__(self, alpha, tau, N ):
		Epoch.__init__(self, alpha, tau, N, False)

	def getMSString( self ):

		#float equality for zero
		if self.tau < 1e-12:
			return (' -G %f' % (self.alpha) )
		else:
			return(' -eG %f %f' % (self.tau, self.alpha) )



def main():
	#test
	p = Population()
	print p
	print
	#
	d = Demography( 10000 )
	d.addEpoch( 1000, 200, True )
	d.addEpoch( -0.5, 15000, False)
	print d
	
	print
	print d.getMSString()
	
	print
	print p.getMSString_Null() + d.getMSString()
	
	ep = d.epochs[0]
	ep.setIterationParameters('alpha', 0, 1000, 10000)
	
	for epItr in ep:
		print d.getMSString()

